/**
 * @author Gao Jian
 * @email gaoj@njupt.edu.cn
 * @create date 2021-06-03 10:38:11
 * @desc 填充空洞或删除小块
 */

#include "defs.h"

__global__ void init_regions(unsigned int *regions, unsigned char *data,
                             unsigned char value, unsigned int nw,
                             unsigned int nh);
__global__ void flip_regions(unsigned char *active, unsigned int *regions,
                             unsigned char mode, unsigned int nw,
                             unsigned int nh);
__global__ void area_regions(unsigned int *areas, unsigned int *regions,
                             unsigned int nw, unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (regions[idx] == 0)
      continue;

    atomicAdd(areas + regions[idx] - 1, 1);
  }
}

__global__ void kcleanparts(unsigned char *data, unsigned int *areas,
                            unsigned int *regions, unsigned char value,
                            unsigned int threshold, unsigned int nw,
                            unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (regions[idx] == 0)
      continue;

    if (areas[regions[idx] - 1] < threshold)
      data[idx] = value;
  }
}

int clean_parts(unsigned char *data, unsigned char val1, unsigned char val2,
                unsigned int threshold, unsigned int nw, unsigned int nh,
                const unsigned verbose = 0) {
  unsigned int nz = nw * nh;
  unsigned char *d_data = NULL;
  unsigned int *d_regions = NULL;
  unsigned int *d_areas = NULL;

  cudaMalloc((void **)&d_data, nz * sizeof(unsigned char));
  cudaMemcpy(d_data, data, nz * sizeof(unsigned char), cudaMemcpyHostToDevice);

  cudaMalloc((void **)&d_regions, nz * sizeof(unsigned int));
  cudaMalloc((void **)&d_areas, nz * sizeof(unsigned int));
  init_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_regions, d_data, val1, nw, nh);

  unsigned char *ctmp = (unsigned char *)d_areas;

  unsigned char active;
  unsigned nit = 0;
  do {
    active = 0;
    cudaMemcpy(ctmp, &active, sizeof(unsigned char), cudaMemcpyHostToDevice);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 0, nw, nh);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 2, nw, nh);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 1, nw, nh);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 3, nw, nh);
    cudaMemcpy(&active, ctmp, sizeof(unsigned char), cudaMemcpyDeviceToHost);
    nit++;
    if (verbose)
      printf("+");
    if (nit % 50 == 0 && verbose)
      printf("  [%3d]\n  ", nit);
  } while (active);
  if (verbose)
    printf("  [%3d]\n", nit);

  nppsSet_32u(0, d_areas, nz);

  area_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_areas, d_regions, nw, nh);
  kcleanparts<<<NUM_BLOCKS, NUM_THREADS>>>(d_data, d_areas, d_regions, val2,
                                           threshold, nw, nh);

  cudaMemcpy(data, d_data, nz * sizeof(unsigned char), cudaMemcpyDeviceToHost);

  cudaFree(d_areas);
  cudaFree(d_regions);
  cudaFree(d_data);

  return 0;
}